package com.lianjia.crawler;

import com.lianjia.crawler.model.HouseInfo;
import com.lianjia.crawler.parser.HouseParser;
import com.lianjia.crawler.utils.DataLoader;
import com.lianjia.crawler.utils.FileUtils;
import com.lianjia.crawler.utils.HttpUtils;
import com.lianjia.crawler.utils.MongoDBUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Scanner;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.Random;

/**
 * 链家爬虫MongoDB版，用于爬取全国城市的二手房信息并存储到MongoDB
 * 北上广杭四个城市各爬取10-20页，其他城市各爬取10-15页
 */
public class MongoDBCrawler {
    private static final Logger logger = LoggerFactory.getLogger(MongoDBCrawler.class);
    private static final int MAX_RETRY = 10000; // 最大重试次数
    private static final int START_PAGE = 10; // 开始页码
    private static final int MAJOR_CITY_PAGES = 20; // 主要城市结束页码
    private static final int OTHER_CITY_PAGES = 15; // 其他城市结束页码
    private static final int THREAD_COUNT = 8; // 线程数量
    
    // 页面之间的等待时间范围（毫秒）
    private static final int MIN_PAGE_WAIT_TIME = 30000; // 最小等待时间：30秒
    private static final int MAX_PAGE_WAIT_TIME = 90000; // 最大等待时间：90秒
    
    // 北上广杭的城市代码
    private static final List<String> MAJOR_CITIES = List.of("北京", "上海", "广州", "杭州");
    
    // MongoDB连接字符串
    private static final String MONGODB_URI = "mongodb://localhost:27017";
    
    // 是否下载图片
    private static boolean downloadImages = false;
    
    // 随机数生成器
    private static final Random random = new Random();
    
    public static void main(String[] args) {
        try {
            logger.info("链家爬虫MongoDB版启动");
            System.out.println("链家爬虫MongoDB版启动");
            
            // 加载城市代码和用户代理
            DataLoader.loadCityCodes();
            DataLoader.loadUserAgents();
            
            // 获取用户输入
            Scanner scanner = new Scanner(System.in);
            
            // 询问是否下载图片
            System.out.print("是否下载房源图片？(y/n，默认n): ");
            String downloadImagesInput = scanner.nextLine().trim().toLowerCase();
            downloadImages = "y".equals(downloadImagesInput) || "yes".equals(downloadImagesInput);
            
            // 询问是否重建MongoDB集合
            System.out.print("是否重建MongoDB集合？(y/n，默认n): ");
            String recreateCollectionInput = scanner.nextLine().trim().toLowerCase();
            boolean recreateCollection = "y".equals(recreateCollectionInput) || "yes".equals(recreateCollectionInput);
            
            // 获取所有城市代码
            Map<String, String> allCities = DataLoader.getAllCityCodes();
            if (allCities == null || allCities.isEmpty()) {
                logger.error("未能加载城市代码，请检查数据文件");
                System.out.println("未能加载城市代码，请检查数据文件");
                return;
            }
            
            // 初始化MongoDB连接
            MongoDBUtils.init(MONGODB_URI, recreateCollection);
            
            // 创建线程池
            ExecutorService executor = Executors.newFixedThreadPool(THREAD_COUNT);
            
            // 统计爬取数量
            AtomicInteger totalCount = new AtomicInteger(0);
            AtomicInteger cityCount = new AtomicInteger(0);
            int totalCities = allCities.size();
            
            // 为每个城市创建爬虫任务
            for (Map.Entry<String, String> city : allCities.entrySet()) {
                String cityName = city.getKey();
                String cityCode = city.getValue();
                
                // 确定爬取页数
                int endPage = MAJOR_CITIES.contains(cityName) ? MAJOR_CITY_PAGES : OTHER_CITY_PAGES;
                
                executor.submit(() -> {
                    try {
                        logger.info("开始爬取城市 {} ({}/{})", cityName, cityCount.incrementAndGet(), totalCities);
                        System.out.println("开始爬取城市 " + cityName + " (" + cityCount.get() + "/" + totalCities + ")");
                        
                        int count = crawlCity(cityName, cityCode, START_PAGE, endPage);
                        totalCount.addAndGet(count);
                        
                        logger.info("城市 {} 爬取完成，共 {} 条数据", cityName, count);
                        System.out.println("城市 " + cityName + " 爬取完成，共 " + count + " 条数据");
                    } catch (Exception e) {
                        logger.error("爬取城市 {} 出错: {}", cityName, e.getMessage(), e);
                        System.out.println("爬取城市 " + cityName + " 出错: " + e.getMessage());
                    }
                });
            }
            
            // 关闭线程池并等待所有任务完成
            executor.shutdown();
            boolean terminated = executor.awaitTermination(48, TimeUnit.HOURS);
            
            if (!terminated) {
                logger.warn("爬虫任务超时，强制关闭线程池");
                System.out.println("爬虫任务超时，强制关闭线程池");
                executor.shutdownNow();
            }
            
            // 关闭MongoDB连接
            MongoDBUtils.close();
            
            logger.info("爬虫任务完成，共爬取 {} 条房源信息", totalCount.get());
            System.out.println("爬虫任务完成，共爬取 " + totalCount.get() + " 条房源信息");
            
        } catch (Exception e) {
            logger.error("爬虫运行出错", e);
            System.out.println("爬虫运行出错: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    /**
     * 爬取指定城市的房源信息
     * 
     * @param cityName 城市名称
     * @param cityCode 城市代码
     * @param startPage 开始页码
     * @param endPage 结束页码
     * @return 爬取的房源数量
     */
    private static int crawlCity(String cityName, String cityCode, int startPage, int endPage) {
        logger.info("开始爬取城市: {}, 页数范围: {}-{}", cityName, startPage, endPage);
        System.out.println("开始爬取城市: " + cityName + ", 页数范围: " + startPage + "-" + endPage);
        
        int totalCount = 0;
        
        for (int page = startPage; page <= endPage; page++) {
            // 构建页面URL
            String url = String.format("https://%s.lianjia.com/ershoufang/pg%d/", cityCode, page);
            String referer = String.format("https://%s.lianjia.com/ershoufang/", cityCode);
            
            // 获取随机用户代理
            String userAgent = DataLoader.getRandomUserAgent();
            
            // 获取页面HTML
            logger.info("正在爬取 {} 第 {} 页...", cityName, page);
            System.out.println("正在爬取 " + cityName + " 第 " + page + " 页...");
            List<HouseInfo> pageInfos = null;
            
            // 不断尝试，直到成功解析出数据或达到最大重试次数
            int retryCount = 0;
            while (pageInfos == null && retryCount < MAX_RETRY) {
                String html = HttpUtils.getHtml(url, userAgent, referer);
                
                // 每次请求使用新的User-Agent
                userAgent = DataLoader.getRandomUserAgent();
                
                if (html != null) {
                    pageInfos = HouseParser.parseHouseList(html);
                }
                
                if (pageInfos == null) {
                    retryCount++;
                    logger.warn("{} 第 {} 页第 {} 次尝试解析失败，正在重试...", cityName, page, retryCount);
                    System.out.println(cityName + " 第 " + page + " 页第 " + retryCount + " 次尝试解析失败，正在重试...");
                    
                    // 如果多次尝试失败，可能是被反爬机制拦截，等待更长时间
                    if (retryCount >= 5) {
                        logger.warn("多次尝试失败，可能是被反爬机制拦截，等待时间增加");
                        System.out.println("多次尝试失败，可能是被反爬机制拦截，等待时间增加");
                        try {
                            Thread.sleep(10000 + retryCount * 2000);
                        } catch (InterruptedException e) {
                            Thread.currentThread().interrupt();
                        }
                    }
                }
            }
            
            if (pageInfos == null) {
                logger.error("{} 第 {} 页爬取失败，已达到最大重试次数", cityName, page);
                System.out.println(cityName + " 第 " + page + " 页爬取失败，已达到最大重试次数");
                continue; // 跳过这一页，继续爬取下一页
            }
            
            // 下载并保存图片（如果需要）
            if (downloadImages) {
                logger.info("正在保存 {} 第 {} 页的图片...", cityName, page);
                System.out.println("正在保存 " + cityName + " 第 " + page + " 页的图片...");
                for (HouseInfo info : pageInfos) {
                    String imageUrl = info.getImage();
                    if (imageUrl != null && !imageUrl.trim().isEmpty()) {
                        String localImagePath = FileUtils.saveImage(imageUrl, cityCode, userAgent);
                        info.setImage(localImagePath);
                    }
                }
            }
            
            // 保存到MongoDB
            int savedCount = MongoDBUtils.batchSaveHouseInfo(pageInfos, cityName);
            totalCount += savedCount;
            
            logger.info("{} 第 {} 页爬取完成，成功保存 {} 条数据，当前共有 {} 条数据", 
                    cityName, page, savedCount, totalCount);
            System.out.println(cityName + " 第 " + page + " 页爬取完成，成功保存 " + savedCount + 
                    " 条数据，当前共有 " + totalCount + " 条数据");
            
            // 如果不是最后一页，则随机等待30-90秒
            if (page < endPage) {
                int waitTime = MIN_PAGE_WAIT_TIME + random.nextInt(MAX_PAGE_WAIT_TIME - MIN_PAGE_WAIT_TIME);
                logger.info("随机等待 {} 秒后继续爬取下一页...", waitTime / 1000);
                System.out.println("随机等待 " + (waitTime / 1000) + " 秒后继续爬取下一页...");
                try {
                    Thread.sleep(waitTime);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            }
        }
        
        logger.info("城市 {} 爬取完成，共获取 {} 条房源信息", cityName, totalCount);
        System.out.println("城市 " + cityName + " 爬取完成，共获取 " + totalCount + " 条房源信息");
        
        return totalCount;
    }
} 